package org.apache.hadoop.hdfs.server.datanode;

import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_DATANODE_SHARED_FILE_DESCRIPTOR_PATHS;
import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_DATANODE_SHARED_FILE_DESCRIPTOR_PATHS_DEFAULT;
import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_SHORT_CIRCUIT_SHARED_MEMORY_WATCHER_INTERRUPT_CHECK_MS;
import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_SHORT_CIRCUIT_SHARED_MEMORY_WATCHER_INTERRUPT_CHECK_MS_DEFAULT;

import java.io.Closeable;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

import com.google.common.annotations.VisibleForTesting;
import org.apache.commons.io.IOUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.InvalidRequestException;
import org.apache.hadoop.hdfs.ExtendedBlockId;
import org.apache.hadoop.hdfs.shortcircuit.ShortCircuitShm;
import org.apache.hadoop.hdfs.shortcircuit.ShortCircuitShm.ShmId;
import org.apache.hadoop.hdfs.shortcircuit.ShortCircuitShm.Slot;
import org.apache.hadoop.hdfs.shortcircuit.ShortCircuitShm.SlotId;
import org.apache.hadoop.io.nativeio.SharedFileDescriptorFactory;
import org.apache.hadoop.net.unix.DomainSocket;
import org.apache.hadoop.net.unix.DomainSocketWatcher;

import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.collect.HashMultimap;

/*
 * ShortCircuitRegistry 类用于管理Datanode侧的所有共享内存，包括创建、释放共享内存等操作。
 * ShortCircuitRegistry定义了segments字段保存所有的共享内存对象RegisteredShm，同时定义了slots字段保存所有的槽位对象Slot
 */
public class ShortCircuitRegistry {
  public static final Log LOG = LogFactory.getLog(ShortCircuitRegistry.class);

  private static final int SHM_LENGTH = 8192;

  /**
   * DfsClientShm类用于描述DFSClient侧的共享内存，与DfsClientShm一样，RegisteredShm也是ShortCircuitRegistry的子类，同样实现了DomainSocketWatcher.Handler接口。
   */
  public static class RegisteredShm extends ShortCircuitShm implements DomainSocketWatcher.Handler {
    private final String clientName;
    private final ShortCircuitRegistry registry;

    RegisteredShm(String clientName, ShmId shmId, FileInputStream stream, ShortCircuitRegistry registry) throws IOException {
      super(shmId, stream);
      this.clientName = clientName;
      this.registry = registry;
    }

    /**
     * Client与Datanode之间通信的domainSocket出现异常时，会调用RegisteredShm.handle()方法处理这个异常。
     * RegisteredShm.handle()方法的逻辑很简单，直接调用ShortCircuitRegistry.removeShm()方法从ShortCircuitRegistry删除这段共享内存即可。
     */
    @Override
    public boolean handle(DomainSocket sock) {
      synchronized (registry) {
        synchronized (this) {
          // 调用ShortCircuitRegistry.removeShm()方法 处理 domainSocket关闭的情况
          registry.removeShm(this);
        }
      }
      return true;
    }

    String getClientName() {
      return clientName;
    }
  }

  /**
   * 用于从ShortCircuitRegistry类中删除一段共享内存，这个方法是在Datanode与DFSClient之间的DomainSocket出现异常时，由RegisteredShm.handle()方法调用的。
   *
   * removeShm()方法首先会停止ShortCircuitRegistry对这段共享内存的追踪，
   * 然后停止对共享内存中所有Slot对象的追踪，
   * 最后释放这段共享内存并关闭共享内存映射文件。
   */
  public synchronized void removeShm(ShortCircuitShm shm) {
    // 停止ShortCircuitRegistry对这段共享内存的追踪
    RegisteredShm removedShm = segments.remove(shm.getShmId());
    Preconditions.checkState(removedShm == shm, "failed to remove " + shm.getShmId());
    // 停止对共享内存中所有Slot对象的追踪
    for (Iterator<Slot> iter = shm.slotIterator(); iter.hasNext(); ) {
      Slot slot = iter.next();
      boolean removed = slots.remove(slot.getBlockId(), slot);
      Preconditions.checkState(removed);
      slot.makeInvalid();
    }
    // 释放这段共享内存并关闭共享内存映射文件
    shm.free();
  }

  /**
   * Whether or not the registry is enabled.
   */
  private boolean enabled;

  /**
   * The factory which creates shared file descriptors.
   */
  private final SharedFileDescriptorFactory shmFactory;
  
  /**
   * A watcher which sends out callbacks when the UNIX domain socket
   * associated with a shared memory segment closes.
   */
  private final DomainSocketWatcher watcher;

  private final HashMap<ShmId, RegisteredShm> segments = new HashMap<ShmId, RegisteredShm>(0);
  
  private final HashMultimap<ExtendedBlockId, Slot> slots = HashMultimap.create(0, 1);
  
  public ShortCircuitRegistry(Configuration conf) throws IOException {
    boolean enabled = false;
    SharedFileDescriptorFactory shmFactory = null;
    DomainSocketWatcher watcher = null;
    try {
      int interruptCheck = conf.getInt(
          DFS_SHORT_CIRCUIT_SHARED_MEMORY_WATCHER_INTERRUPT_CHECK_MS,
          DFS_SHORT_CIRCUIT_SHARED_MEMORY_WATCHER_INTERRUPT_CHECK_MS_DEFAULT);
      if (interruptCheck <= 0) {
        throw new IOException( DFS_SHORT_CIRCUIT_SHARED_MEMORY_WATCHER_INTERRUPT_CHECK_MS + " was set to " + interruptCheck);
      }
      String shmPaths[] = conf.getTrimmedStrings(DFS_DATANODE_SHARED_FILE_DESCRIPTOR_PATHS);
      if (shmPaths.length == 0) {
        shmPaths = DFS_DATANODE_SHARED_FILE_DESCRIPTOR_PATHS_DEFAULT.split(",");
      }
      shmFactory = SharedFileDescriptorFactory.create("HadoopShortCircuitShm_", shmPaths);
      String dswLoadingFailure = DomainSocketWatcher.getLoadingFailureReason();
      if (dswLoadingFailure != null) {
        throw new IOException(dswLoadingFailure);
      }
      watcher = new DomainSocketWatcher(interruptCheck);
      enabled = true;
    } catch (IOException e) {
    } finally {
      this.enabled = enabled;
      this.shmFactory = shmFactory;
      this.watcher = watcher;
    }
  }

  /**
   * 于在缓存数据块对应的Slot对象上添加可锚状态位，这个方法是在Datanode将一个数据块添加到缓存时，由CachingTask对象调用的。
   *
   * processBlockMlockEvent()会从ShortCircuit Registry.slots字段获取该数据块对应的Slot对象，
   * 然后调用Slot.makeAnchorable()方法添加一个可锚状态位，由于这个状态位是保存在共享内存文件中的，所以这个状态位的信息会同步到DFSClient端的Slot对象上。
   */
  public synchronized void processBlockMlockEvent(ExtendedBlockId blockId) {
    if (!enabled) {
      return;
    }
    Set<Slot> affectedSlots = slots.get(blockId);
    for (Slot slot : affectedSlots) {
      slot.makeAnchorable();
    }
  }

  /**
   * 类似于processBlockMlockEvent()方法，当Datanode将数据块从缓存中移除时，会调用这个方法删除Slot对象的可锚状态位。
   */
  public synchronized boolean processBlockMunlockRequest( ExtendedBlockId blockId) {
    if (!enabled) {
      return true;
    }
    boolean allowMunlock = true;
    Set<Slot> affectedSlots = slots.get(blockId);
    for (Slot slot : affectedSlots) {
      slot.makeUnanchorable();
      if (slot.isAnchored()) {
        allowMunlock = false;
      }
    }
    return allowMunlock;
  }

  /**
   * Invalidate any slot associated with a blockId that we are invalidating
   * (deleting) from this DataNode.  When a slot is invalid, the DFSClient will
   * not use the corresponding replica for new read or mmap operations (although
   * existing, ongoing read or mmap operations will complete.)
   *
   * @param blockId        The block ID.
   */
  public synchronized void processBlockInvalidation(ExtendedBlockId blockId) {
    if (!enabled) {
      return;
    }
    final Set<Slot> affectedSlots = slots.get(blockId);
    if (!affectedSlots.isEmpty()) {
      final StringBuilder bld = new StringBuilder();
      String prefix = "";
      bld.append("Block ").append(blockId).append(" has been invalidated.  ").
          append("Marking short-circuit slots as invalid: ");
      for (Slot slot : affectedSlots) {
        slot.makeInvalid();
        bld.append(prefix).append(slot.toString());
        prefix = ", ";
      }
      LOG.info(bld.toString());
    }
  }

  public synchronized String getClientNames(ExtendedBlockId blockId) {
    if (!enabled) return "";
    final HashSet<String> clientNames = new HashSet<String>();
    final Set<Slot> affectedSlots = slots.get(blockId);
    for (Slot slot : affectedSlots) {
      clientNames.add(((RegisteredShm)slot.getShm()).getClientName());
    }
    return Joiner.on(",").join(clientNames);
  }

  public static class NewShmInfo implements Closeable {
    public final ShmId shmId;
    public final FileInputStream stream;

    NewShmInfo(ShmId shmId, FileInputStream stream) {
      this.shmId = shmId;
      this.stream = stream;
    }

    @Override
    public void close() throws IOException {
      stream.close();
    }
  }

  /**
   * 用于在Datanode侧构造一段共享内存，
   * 这个方法是在DataXceiver.requestShortCircuitShm()响应DataTransferProtocol.requestShortCircuitShm()请求时调用的。
   *
   * createNewMemorySegment()方法
   * 1. 打开共享内存文件，并将该文件映射到内存中。
   * 2. 然后创建RegisteredShm对象管理该共享内存，并将RegisteredShm对象加入ShortCircuitRegistry.segments字段中保存。
   */
  public NewShmInfo createNewMemorySegment(String clientName, DomainSocket sock) throws IOException {
    NewShmInfo info = null;
    RegisteredShm shm = null;
    ShmId shmId = null;
    synchronized (this) {
      if (!enabled) {
        throw new UnsupportedOperationException();
      }
      FileInputStream fis = null;
      try {
        do {
          shmId = ShmId.createRandom();
        } while (segments.containsKey(shmId));
        fis = shmFactory.createDescriptor(clientName, SHM_LENGTH); // TODO
        shm = new RegisteredShm(clientName, shmId, fis, this);
      } finally {
        if (shm == null) {
          IOUtils.closeQuietly(fis);
        }
      }
      info = new NewShmInfo(shmId, fis);
      segments.put(shmId, shm);
    }
    // Drop the registry lock to prevent deadlock.
    // After this point, RegisteredShm#handle may be called at any time.
    watcher.add(sock, shm);
    return info;
  }

  /**
   * 用于在Datanode侧的共享内存中添加一个Slot对象，这个方法是在DataXceiver.requestShortCircuitFds()响应DataTransferProtocol.requestShortCircuitFds()请求时调用的。
   *
   * registerSlot()方法首先从ShortCircuitRegistry.segments字段中取出指定共享内存对应的RegisteredShm对象，
   * 然后调用RegisteredShm.registerSlot()方法构造并添加Slot对象，
   * 最后将该Slot对象加入ShortCircuitRegistry.slots字段中保存。
   */
  public synchronized void registerSlot(ExtendedBlockId blockId, SlotId slotId, boolean isCached) throws InvalidRequestException {
    if (!enabled) {
      throw new UnsupportedOperationException();
    }
    ShmId shmId = slotId.getShmId();
    RegisteredShm shm = segments.get(shmId);
    if (shm == null) {
      throw new InvalidRequestException("there is no shared memory segment registered with shmId " + shmId);
    }
    Slot slot = shm.registerSlot(slotId.getSlotIdx(), blockId);
    if (isCached) {
      slot.makeAnchorable();
    } else {
      slot.makeUnanchorable();
    }
    boolean added = slots.put(blockId, slot);
    Preconditions.checkState(added);
  }

  /**
   * 用于在Datanode侧的共享内存中删除一个Slot对象，这个方法是在DataXceiver.releaseShortCircuitFds()响应DataTransferProtocol.releaseShortCircuitFds()请求时调用的。
   *
   * unregisterSlot()方法首先从ShortCircuitRegistry.segments字段中取出指定共享内存对应的RegisteredShm对象，
   * 然后调用RegisteredShm.unregisterSlot()释放Slot对象，
   * 最后将该Slot对象从ShortCircuitRegistry.slots字段中删除。
   */
  public synchronized void unregisterSlot(SlotId slotId) throws InvalidRequestException {
    if (!enabled) {
      throw new UnsupportedOperationException();
    }
    ShmId shmId = slotId.getShmId();
    RegisteredShm shm = segments.get(shmId);
    if (shm == null) {
      throw new InvalidRequestException("there is no shared memory segment " + "registered with shmId " + shmId);
    }
    Slot slot = shm.getSlot(slotId.getSlotIdx());
    slot.makeInvalid();
    shm.unregisterSlot(slotId.getSlotIdx());
    slots.remove(slot.getBlockId(), slot);
  }
  
  public void shutdown() {
    synchronized (this) {
      if (!enabled) {
        return;
      }
      enabled = false;
    }
    IOUtils.closeQuietly(watcher);
  }

  public static interface Visitor {
    void accept(HashMap<ShmId, RegisteredShm> segments,
                HashMultimap<ExtendedBlockId, Slot> slots);
  }

  @VisibleForTesting
  public synchronized void visit(Visitor visitor) {
    visitor.accept(segments, slots);
  }
}
